# 1.使用sklearn配合pytorch完成以下操作（每题10分）
# (1)数据处理
import matplotlib.pyplot as plt
import torch
from sklearn.datasets import load_iris

torch.manual_seed(1234)

# ①读取iris数据集
x, y = load_iris(return_X_y=True)
# ②自定义超参数数值
# ③将数据7:3切分
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(x, y, train_size=0.7)
X_train = torch.tensor(X_train)
X_test = torch.tensor(X_test)
y_train = torch.tensor(y_train)
y_test = torch.tensor(y_test)
# ④放入dataloder
# ⑤创建全连接层，两个隐藏层，神经元数量分别是10,5
# ⑥每个隐藏层后用relu激活，dropout设定0.2

model = torch.nn.Sequential(
    torch.nn.Linear(in_features=4, out_features=10),
    torch.nn.ReLU(),
    torch.nn.Dropout(0.2),
    torch.nn.Linear(in_features=10, out_features=5),
    torch.nn.ReLU(),
    torch.nn.Dropout(0.2),
    torch.nn.Linear(in_features=5, out_features=3),
)
# ⑦使用交叉熵作为代价函数
loss_fn = torch.nn.CrossEntropyLoss()
op = torch.optim.Adagrad(model.parameters(), lr=0.1)
print()
loss_list = []
# ⑧训练模型，每100批次打印损失值和准确率
for i in range(10000):
    op.zero_grad()
    predict = model(X_train.float())
    loss = loss_fn(predict, y_train)
    loss.backward()
    op.step()
    loss_list.append(loss.item())
    print(i, loss.item())
# ⑨绘制损失变化曲线
plt.plot(loss_list)
plt.show()
# ⑩打印最终准确率
avg_acc = 0
with torch.no_grad():
    predict = model(X_test.float())
    avg_acc += predict.argmax(1).eq(y_test).int().sum().item()
print('acc is', avg_acc / len(y_test))
